#include "leif.h"
#define GLFW_EXPOSE_NATIVE_X11
#include <GLFW/glfw3.h>
#include <GLFW/glfw3native.h>
#include <X11/Xatom.h>
#include <X11/Xlib.h>
#include <cglm/types-struct.h>
#include <security/pam_appl.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <time.h>
#include <unistd.h>

#define DEBUG_MODE 0
#define FONT_PATH "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"
#define MAIN_FONT_SIZE 28
#define HINT_FONT_SIZE 20
#define MAX_PASSWORD_LEN 512
#define MAX_FAILED_ATTEMPTS 5
#define TARGET_FPS 60
#define FRAME_TIME_US (1000000 / TARGET_FPS)
#define POPUP_WIDTH 800.0f
#define POPUP_HEIGHT 300.0f

#define LF_COLOR_BG (LfColor){0, 0, 0, 255}
#define LF_COLOR_INPUT_BG (LfColor){40, 40, 40, 255}
#define LF_COLOR_TEXT (LfColor){255, 255, 255, 255}
#define LF_COLOR_PROMPT (LfColor){255, 204, 0, 255}
#define LF_COLOR_ERROR (LfColor){255, 0, 0, 255}
#define LF_COLOR_SUCCESS (LfColor){0, 255, 0, 255}
#define LF_COLOR_SEPARATOR (LfColor){80, 80, 80, 255}
#define LF_COLOR_HINT (LfColor){138, 138, 138, 255}

char password[MAX_PASSWORD_LEN] = {0};
size_t cursor = 0;
int input_len = 0;
int failed_attempts = 0;
bool is_locked = true;
bool show_help = false;
bool redraw_needed = true;
bool cursor_visible = true;
double last_cursor_flip = 0.0;
double prev_mouse_x = -1, prev_mouse_y = -1;

GLFWwindow *glfw_window;
Display *x11_dpy;
LfFont main_font;
LfFont hint_font;
LfTheme app_theme;
static LfInputField password_input = {0};
char pass_indicator[MAX_PASSWORD_LEN + 1] = {0};

static int pam_conv_callback(int num_msg, const struct pam_message **msg,
                             struct pam_response **resp, void *appdata_ptr) {
  (void)appdata_ptr;
  *resp = calloc(num_msg, sizeof(struct pam_response));
  if (!*resp)
    return PAM_BUF_ERR;

  for (int i = 0; i < num_msg; i++) {
    if (msg[i]->msg_style == PAM_PROMPT_ECHO_OFF) {
      resp[i]->resp = strdup(password);
      resp[i]->resp_retcode = 0;
    }
  }
  return PAM_SUCCESS;
}

static bool verify_password() {
  if (input_len == 0)
    return false;

#ifdef DEBUG_MODE
  printf("[DEBUG] Skip PAM verification, accept input\n");
  return true;
#else
  struct pam_conv conv = {pam_conv_callback, NULL};
  pam_handle_t *pam_handle = NULL;
  const char *username = getlogin();
  int ret = pam_start("login", username, &conv, &pam_handle);

  if (ret != PAM_SUCCESS) {
    fprintf(stderr, "PAM init failed: %s\n", pam_strerror(pam_handle, ret));
    pam_end(pam_handle, ret);
    return false;
  }

  ret = pam_authenticate(pam_handle, 0);
  pam_end(pam_handle, ret);

  if (ret == PAM_SUCCESS)
    return true;
  failed_attempts++;
  return false;
#endif
}

static void clear_password() {
  memset(password, 0, sizeof(password));
  memset(pass_indicator, 0, sizeof(pass_indicator));
  memset(password_input.buf, 0, password_input.buf_size);
  input_len = 0;
  cursor = 0;
  password_input.cursor_index = 0; // 仅同步光标，删除len
  redraw_needed = true;
}

static bool grab_x11_input(Window win) {
  if (XGrabKeyboard(x11_dpy, win, False, GrabModeAsync, GrabModeAsync,
                    CurrentTime) != GrabSuccess) {
    fprintf(stderr, "Failed to grab keyboard (retrying...)\n");
    usleep(100000);
    if (XGrabKeyboard(x11_dpy, win, False, GrabModeAsync, GrabModeAsync,
                      CurrentTime) != GrabSuccess) {
      fprintf(stderr, "Fatal: Could not grab keyboard\n");
      return false;
    }
  }

  if (XGrabPointer(x11_dpy, win, False,
                   ButtonPressMask | ButtonReleaseMask | PointerMotionMask,
                   GrabModeAsync, GrabModeAsync, None, None,
                   CurrentTime) != GrabSuccess) {
    fprintf(stderr, "Failed to grab pointer (retrying...)\n");
    usleep(100000);
    if (XGrabPointer(x11_dpy, win, False,
                     ButtonPressMask | ButtonReleaseMask | PointerMotionMask,
                     GrabModeAsync, GrabModeAsync, None, None,
                     CurrentTime) != GrabSuccess) {
      fprintf(stderr, "Fatal: Could not grab pointer\n");
      XUngrabKeyboard(x11_dpy, CurrentTime);
      return false;
    }
  }

  return true;
}

static bool init_lf_resources() {
  x11_dpy = XOpenDisplay(NULL);
  if (!x11_dpy) {
    fprintf(stderr, "Failed to open X display\n");
    return false;
  }

  if (!glfwInit()) {
    fprintf(stderr, "GLFW Init Failed\n");
    XCloseDisplay(x11_dpy);
    return false;
  }

  GLFWmonitor *monitor = glfwGetPrimaryMonitor();
  const GLFWvidmode *vidmode = glfwGetVideoMode(monitor);
  glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
  glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
  glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
  glfwWindowHint(GLFW_OPENGL_FORWARD_COMPAT, GLFW_TRUE);
  glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);
  glfwWindowHint(GLFW_DECORATED, GLFW_FALSE);

  glfw_window = glfwCreateWindow(vidmode->width, vidmode->height, "lf-xlock",
                                 monitor, NULL);
  if (!glfw_window) {
    fprintf(stderr, "Fullscreen Window Creation Failed\n");
    glfwTerminate();
    XCloseDisplay(x11_dpy);
    return false;
  }

  Window glfw_win_xid = glfwGetX11Window(glfw_window);
  Atom net_wm_state = XInternAtom(x11_dpy, "_NET_WM_STATE", False);
  Atom net_wm_state_fullscreen =
      XInternAtom(x11_dpy, "_NET_WM_STATE_FULLSCREEN", False);
  XChangeProperty(x11_dpy, glfw_win_xid, net_wm_state, XA_ATOM, 32,
                  PropModeReplace, (unsigned char *)&net_wm_state_fullscreen,
                  1);

#ifndef DEBUG_MODE
  if (!grab_x11_input(glfw_win_xid)) {
    glfwDestroyWindow(glfw_window);
    glfwTerminate();
    XCloseDisplay(x11_dpy);
    return false;
  }
#else
  printf("[DEBUG] Skip X11 input grab\n");
#endif

  glfwMakeContextCurrent(glfw_window);
  glfwSetInputMode(glfw_window, GLFW_CURSOR, GLFW_CURSOR_NORMAL);
  glfwSwapInterval(1);

  lf_init_glfw(vidmode->width, vidmode->height, glfw_window);

  main_font = lf_load_font(FONT_PATH, MAIN_FONT_SIZE);
  if (main_font.font_info == NULL) {
    fprintf(stderr, "Main Font Load Failed: %s\n", FONT_PATH);
    lf_terminate();
    glfwDestroyWindow(glfw_window);
    glfwTerminate();
    XCloseDisplay(x11_dpy);
    return false;
  }

  hint_font = lf_load_font(FONT_PATH, HINT_FONT_SIZE);
  if (hint_font.font_info == NULL) {
    fprintf(stderr, "Hint Font Load Failed: %s\n", FONT_PATH);
    lf_free_font(&main_font);
    lf_terminate();
    glfwDestroyWindow(glfw_window);
    glfwTerminate();
    XCloseDisplay(x11_dpy);
    return false;
  }

  app_theme = lf_default_theme();
  app_theme.div_props.color = LF_COLOR_BG;
  app_theme.div_props.padding = 0.0f;
  app_theme.text_props.text_color = LF_COLOR_TEXT;
  app_theme.inputfield_props.color = LF_COLOR_INPUT_BG;
  app_theme.inputfield_props.text_color = LF_COLOR_SUCCESS;
  app_theme.inputfield_props.border_width = 0.0f;
  app_theme.inputfield_props.padding = 6.0f;
  app_theme.inputfield_props.corner_radius = 0.0f;
  lf_set_theme(app_theme);

  password_input.buf = pass_indicator;
  password_input.buf_size = sizeof(pass_indicator);
  password_input.placeholder = "";
  password_input.width = POPUP_WIDTH - 40.0f;
  password_input.height = MAIN_FONT_SIZE + 12.0f;
  password_input.max_chars = MAX_PASSWORD_LEN - 1;
  password_input.retain_height = true;
  password_input._init = true;
  // 移除对len的初始化（结构体无此成员）

#ifdef __linux__
  if (mlock(password, sizeof(password)) != 0) {
    fprintf(stderr, "Warning: Could not lock password memory (run with "
                    "sufficient privileges)\n");
  }
#endif

  return true;
}

static void draw_cursor(vec2s pos) {
  vec2s cursor_size = {2.0f, (float)main_font.font_size};
  lf_rect_render(pos, cursor_size, LF_COLOR_TEXT, LF_NO_COLOR, 0.0f, 0.0f);
}

static void render() {
  if (!redraw_needed)
    return;

  lf_begin();

  vec2s screen_size = {(float)lf_get_display_width(),
                       (float)lf_get_display_height()};
  lf_rect(screen_size.x, screen_size.y, LF_COLOR_BG, 0.0f);

  vec2s popup_pos = {(screen_size.x - POPUP_WIDTH) / 2.0f,
                     (screen_size.y - POPUP_HEIGHT) / 2.0f};
  vec2s popup_size = {POPUP_WIDTH, POPUP_HEIGHT};
  lf_div_begin(popup_pos, popup_size, false);

  lf_push_font(&main_font);

  const char *title = "Screen Locked";
#ifdef DEBUG_MODE
  title = "Screen Locked [DEBUG MODE]";
#endif
  vec2s title_dim = lf_text_dimension(title);
  lf_set_ptr_x((POPUP_WIDTH - title_dim.x) / 2.0f);
  lf_set_ptr_y(20.0f);
  lf_text(title);

  float input_y = 20.0f + main_font.font_size + 20.0f;
  vec2s input_box_pos = {20.0f, input_y};
  vec2s input_box_size = {POPUP_WIDTH - 40.0f,
                          (float)main_font.font_size + 12.0f};
  lf_rect_render(input_box_pos, input_box_size, LF_COLOR_INPUT_BG, LF_NO_COLOR,
                 0.0f, 0.0f);

  const char *prompt = "Password: ";
  vec2s prompt_dim = lf_text_dimension(prompt);
  lf_set_ptr_x(28.0f);
  lf_set_ptr_y(input_y + 6.0f);
  LfTheme temp_theme = lf_get_theme();
  temp_theme.text_props.text_color = LF_COLOR_PROMPT;
  lf_set_theme(temp_theme);
  lf_text(prompt);

  lf_set_ptr_x(28.0f + prompt_dim.x);
  lf_set_ptr_y(input_y + 6.0f);
  temp_theme.text_props.text_color = LF_COLOR_SUCCESS;
  temp_theme.inputfield_props.text_color = LF_COLOR_SUCCESS;
  lf_set_theme(temp_theme);
  lf_input_text(&password_input);

  double now = (double)clock() / CLOCKS_PER_SEC;
  if (now - last_cursor_flip > 0.5) {
    cursor_visible = !cursor_visible;
    last_cursor_flip = now;
    redraw_needed = true;
  }
  if (cursor_visible && input_len > 0) {
    char cursor_substr[MAX_PASSWORD_LEN];
    strncpy(cursor_substr, pass_indicator, cursor);
    cursor_substr[cursor] = '\0';
    vec2s cursor_substr_dim = lf_text_dimension(cursor_substr);
    vec2s cursor_pos = {28.0f + prompt_dim.x + cursor_substr_dim.x,
                        input_y + 6.0f};
    draw_cursor(cursor_pos);
  }

  float separator_y = input_y + main_font.font_size + 20.0f;
  vec2s separator_pos = {20.0f, separator_y};
  vec2s separator_size = {POPUP_WIDTH - 40.0f, 1.0f};
  lf_rect_render(separator_pos, separator_size, LF_COLOR_SEPARATOR, LF_NO_COLOR,
                 0.0f, 0.0f);

  if (failed_attempts > 0) {
    char fail_text[64];
    snprintf(fail_text, sizeof(fail_text), "Failed: %d/%d", failed_attempts,
             MAX_FAILED_ATTEMPTS);
    vec2s fail_dim = lf_text_dimension(fail_text);
    lf_set_ptr_x((POPUP_WIDTH - fail_dim.x) / 2.0f);
    lf_set_ptr_y(separator_y + 15.0f);
    temp_theme.text_props.text_color = LF_COLOR_ERROR;
    lf_set_theme(temp_theme);
    lf_text(fail_text);

    if (failed_attempts >= MAX_FAILED_ATTEMPTS) {
      const char *block = "Temporarily Blocked";
      vec2s block_dim = lf_text_dimension(block);
      lf_set_ptr_x((POPUP_WIDTH - block_dim.x) / 2.0f);
      lf_set_ptr_y(separator_y + 15.0f + main_font.font_size);
      lf_text(block);
    }
  }

  lf_pop_font();
  lf_push_font(&hint_font);

  temp_theme.text_props.text_color = LF_COLOR_HINT;
  lf_set_theme(temp_theme);

  if (show_help) {
    const char *help[] = {
        "ESC: Clear | Enter: Submit | F1: Hide Help",
        "Ctrl+A: Home | Ctrl+E: End | Ctrl+F/B: Left/Right",
        "Backspace: Delete | Ctrl+D: Delete Forward | Ctrl+K: Clear Rest"};
    float help_y = separator_y + 40.0f;
    for (int i = 0; i < 3; i++) {
      vec2s help_dim = lf_text_dimension(help[i]);
      lf_set_ptr_x((POPUP_WIDTH - help_dim.x) / 2.0f);
      lf_set_ptr_y(help_y + i * (HINT_FONT_SIZE + 10.0f));
      lf_text(help[i]);
    }
  } else {
    const char *hint = "Press F1 for help";
    vec2s hint_dim = lf_text_dimension(hint);
    lf_set_ptr_x(POPUP_WIDTH - 20.0f - hint_dim.x - 5.0f);
    lf_set_ptr_y(POPUP_HEIGHT - 20.0f - HINT_FONT_SIZE);
    lf_text(hint);
  }

  lf_pop_font();
  lf_set_theme(app_theme);
  lf_div_end();
  lf_end();

  glfwSwapBuffers(glfw_window);
  redraw_needed = false;
}

static void handle_key_press() {
  if (!is_locked)
    return;

  bool ctrl_pressed = lf_key_is_down(GLFW_KEY_LEFT_CONTROL) ||
                      lf_key_is_down(GLFW_KEY_RIGHT_CONTROL);
  bool mod_pressed = lf_key_is_down(GLFW_KEY_LEFT_SUPER) ||
                     lf_key_is_down(GLFW_KEY_RIGHT_SUPER);

  if (mod_pressed)
    return;

  if (ctrl_pressed) {
    if (lf_key_changed(GLFW_KEY_A) && lf_key_went_down(GLFW_KEY_A)) {
      cursor = 0;
      password_input.cursor_index = 0;
      redraw_needed = true;
      return;
    }
    if (lf_key_changed(GLFW_KEY_E) && lf_key_went_down(GLFW_KEY_E)) {
      cursor = input_len;
      password_input.cursor_index = input_len;
      redraw_needed = true;
      return;
    }
    if (lf_key_changed(GLFW_KEY_F) && lf_key_went_down(GLFW_KEY_F)) {
      if (cursor < input_len) {
        cursor++;
        password_input.cursor_index++;
        redraw_needed = true;
      }
      return;
    }
    if (lf_key_changed(GLFW_KEY_B) && lf_key_went_down(GLFW_KEY_B)) {
      if (cursor > 0) {
        cursor--;
        password_input.cursor_index--;
        redraw_needed = true;
      }
      return;
    }
    if (lf_key_changed(GLFW_KEY_D) && lf_key_went_down(GLFW_KEY_D)) {
      if (cursor < input_len) {
        memmove(&password[cursor], &password[cursor + 1],
                MAX_PASSWORD_LEN - cursor - 1);
        memmove(&pass_indicator[cursor], &pass_indicator[cursor + 1],
                MAX_PASSWORD_LEN - cursor);
        input_len--;
        password_input.cursor_index = cursor;
        // 移除对len的赋值（结构体无此成员）
        redraw_needed = true;
      }
      return;
    }
    if (lf_key_changed(GLFW_KEY_K) && lf_key_went_down(GLFW_KEY_K)) {
      password[cursor] = '\0';
      pass_indicator[cursor] = '\0';
      input_len = cursor;
      password_input.cursor_index = cursor;
      // 移除对len的赋值
      redraw_needed = true;
      return;
    }
  }

  if (lf_key_changed(GLFW_KEY_ESCAPE) && lf_key_went_down(GLFW_KEY_ESCAPE)) {
    clear_password();
    return;
  }

  if (lf_key_changed(GLFW_KEY_ENTER) && lf_key_went_down(GLFW_KEY_ENTER)) {
    if (verify_password()) {
      is_locked = false;
      printf("[INFO] Unlock successful\n");
    } else {
      clear_password();
      printf("[INFO] Unlock failed, clear input\n");
    }
    redraw_needed = true;
    return;
  }

  if (lf_key_changed(GLFW_KEY_BACKSPACE) &&
      lf_key_went_down(GLFW_KEY_BACKSPACE)) {
    if (cursor > 0) {
      cursor--;
      memmove(&password[cursor], &password[cursor + 1],
              MAX_PASSWORD_LEN - cursor - 1);
      memmove(&pass_indicator[cursor], &pass_indicator[cursor + 1],
              MAX_PASSWORD_LEN - cursor);
      input_len--;
      password_input.cursor_index = cursor;
      // 移除对len的赋值
      redraw_needed = true;
    }
    return;
  }

  if (lf_key_changed(GLFW_KEY_DELETE) && lf_key_went_down(GLFW_KEY_DELETE)) {
    if (cursor < input_len) {
      memmove(&password[cursor], &password[cursor + 1],
              MAX_PASSWORD_LEN - cursor - 1);
      memmove(&pass_indicator[cursor], &pass_indicator[cursor + 1],
              MAX_PASSWORD_LEN - cursor);
      input_len--;
      password_input.cursor_index = cursor;
      // 移除对len的赋值
      redraw_needed = true;
    }
    return;
  }

  if (lf_key_changed(GLFW_KEY_RIGHT) && lf_key_went_down(GLFW_KEY_RIGHT)) {
    if (cursor < input_len) {
      cursor++;
      password_input.cursor_index++;
      redraw_needed = true;
    }
    return;
  }

  if (lf_key_changed(GLFW_KEY_LEFT) && lf_key_went_down(GLFW_KEY_LEFT)) {
    if (cursor > 0) {
      cursor--;
      password_input.cursor_index--;
      redraw_needed = true;
    }
    return;
  }

  if (lf_key_changed(GLFW_KEY_HOME) && lf_key_went_down(GLFW_KEY_HOME)) {
    cursor = 0;
    password_input.cursor_index = 0;
    redraw_needed = true;
    return;
  }

  if (lf_key_changed(GLFW_KEY_END) && lf_key_went_down(GLFW_KEY_END)) {
    cursor = input_len;
    password_input.cursor_index = input_len;
    redraw_needed = true;
    return;
  }

  if (lf_key_changed(GLFW_KEY_F1) && lf_key_went_down(GLFW_KEY_F1)) {
    show_help = !show_help;
    redraw_needed = true;
    return;
  }

  LfCharEvent char_event = lf_char_event();
  if (char_event.happened && char_event.charcode >= 32 &&
      char_event.charcode <= 126 && input_len < MAX_PASSWORD_LEN - 1) {
    memmove(&password[cursor + 1], &password[cursor],
            MAX_PASSWORD_LEN - cursor - 1);
    memmove(&pass_indicator[cursor + 1], &pass_indicator[cursor],
            MAX_PASSWORD_LEN - cursor);
    password[cursor] = (char)char_event.charcode;
    pass_indicator[cursor] = '*';
    cursor++;
    input_len++;
    password_input.cursor_index = cursor;
    // 移除对len的赋值
    redraw_needed = true;
  }
}

static void handle_events() {
  if (lf_key_event().happened) {
    handle_key_press();
  }

  double mouse_x, mouse_y;
  glfwGetCursorPos(glfw_window, &mouse_x, &mouse_y);
  if ((prev_mouse_x != -1 && prev_mouse_y != -1) &&
      (mouse_x != prev_mouse_x || mouse_y != prev_mouse_y)) {
    redraw_needed = true;
    prev_mouse_x = mouse_x;
    prev_mouse_y = mouse_y;
  } else if (prev_mouse_x == -1 && prev_mouse_y == -1) {
    prev_mouse_x = mouse_x;
    prev_mouse_y = mouse_y;
  }

  for (int btn = 0; btn < GLFW_MOUSE_BUTTON_LAST; btn++) {
    if (glfwGetMouseButton(glfw_window, btn) == GLFW_PRESS) {
      redraw_needed = true;
      break;
    }
  }

  glfwPollEvents();
}

static void cleanup() {
#ifndef DEBUG_MODE
  XUngrabKeyboard(x11_dpy, CurrentTime);
  XUngrabPointer(x11_dpy, CurrentTime);
#endif

  memset(password, 0, sizeof(password));
  memset(pass_indicator, 0, sizeof(pass_indicator));

#ifdef __linux__
  munlock(password, sizeof(password));
#endif

  lf_free_font(&main_font);
  lf_free_font(&hint_font);
  lf_terminate();

  glfwDestroyWindow(glfw_window);
  glfwTerminate();
  XCloseDisplay(x11_dpy);
}

int main() {
  if (!init_lf_resources()) {
    cleanup();
    return 1;
  }

  struct timespec last_frame;
  clock_gettime(CLOCK_MONOTONIC, &last_frame);

  while (is_locked && !glfwWindowShouldClose(glfw_window)) {
    handle_events();
    render();

    struct timespec now;
    clock_gettime(CLOCK_MONOTONIC, &now);
    long elapsed = (now.tv_sec - last_frame.tv_sec) * 1000000 +
                   (now.tv_nsec - last_frame.tv_nsec) / 1000;
    if (elapsed < FRAME_TIME_US) {
      usleep(FRAME_TIME_US - elapsed);
    }
    last_frame = now;
  }

  cleanup();
  return 0;
}